import os
import numpy as np
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
import argparse

# Define bins and labels
bins = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
labels = ['[0,0]', '(0,10]', '(10,20]', '(20,30]', '(30,40]', '(40,50]', 
          '(50,60]', '(60,70]', '(70,80]', '(80,90]', '(90,100)', '[100,100]']

def read_tsv(file_path,cov):
    """Read a TSV file and return a dictionary with the first two columns as keys and the 11th column as values"""
    data = {}
    with open(file_path, 'r') as file:
        for line in file:
            columns = line.strip().split('\t')
            if int(columns[9])<cov:
                continue
            key = (columns[0], columns[1])  # Use the first two columns as the key
            value = float(columns[10])  # Use the 11th column as the value (index 10)
            data[key] = (columns[9],value)
    return data

def assign_keys_to_bins(file1_data, bins, labels):
    """Assign keys to corresponding bins based on file1 data"""
    binned_keys = {label: set() for label in labels}
    for key, (cov,value) in file1_data.items():
        if value == 0:
            binned_keys['[0,0]'].add(key)
        elif value == 100:
            binned_keys['[100,100]'].add(key)
        else:
            for i in range(len(bins) - 1):
                if bins[i] < value <= bins[i + 1]:
                    binned_keys[labels[i + 1]].add(key)
                    break
    return binned_keys

def filter_data_by_keys(data, keys):
    """Filter a data dictionary by a set of keys"""
    return {k: data[k] for k in keys if k in data}

def sum_of_squared_differences(data1, data2, intersection_keys):
    values1 = [data1[k][1]/100.0 for k in intersection_keys]
    values2 = [data2[k][1]/100.0 for k in intersection_keys]
    return sum([(a - b) ** 2 for a, b in zip(values1, values2)])

def calculate_rmse(data1, data2, intersection_keys):
    """Calculate RMSE between two datasets, possibly limited to a set of intersection keys"""
    # if intersection_keys:
    #     common_keys = intersection_keys & set(data1.keys()) & set(data2.keys())
    # else:
    #     common_keys = set(data1.keys()) & set(data2.keys())

    #if common_keys:
    values1 = [data1[k][1]/100.0 for k in intersection_keys]
    values2 = [data2[k][1]/100.0 for k in intersection_keys]
    mse=mean_squared_error(values1, values2)
    return np.sqrt(mse)
    #return None, 0

def plot_bar_chart(values1, values2, bin_label, output_dir,intersection_keys1,intersection_keys2=None):
    """Create and save a bar chart for the given values in the specified bin"""
    plt.figure()
    y1 = [values1[k][1] for k in intersection_keys1]
    if intersection_keys2 is not None:
        y2 = [values2[k][1] for k in intersection_keys2]
    else:
        y2 = [values2[k][1] for k in intersection_keys1]
    #colors = ['r', 'g']
    width = 1#0.35
    heights1 = np.bincount(y1, minlength=40)
    heights2 = np.bincount(y2, minlength=40)
    #heights1=heights1[:40]
    #heights2=heights2[:40]
    fig, ax = plt.subplots()
    ax.bar(np.arange(0, 101), heights1, width,color='r',alpha=0.3, label='dorado')
    ax.bar(np.arange(0, 101), heights2, width,color='g',alpha=0.3, label='T5')
    ax.legend()
    plt.xlabel('prediction value')
    plt.ylabel('Frequency')
    plt.title(f'Predictions in Bin {bin_label}, HG002,R10.4.1,chr1')
    plt.xlim(-1, 101)
    plt.tight_layout()
    os.makedirs(output_dir, exist_ok=True)
    plt.savefig(os.path.join(output_dir, f'bin_{bin_label.replace(",", "_").replace(" ", "")}.png'))
    plt.close()

def plot_histogram(file2_values, output_dir, label):
    """Plot histogram for values from file2 and save to the output directory"""
    # Generate histogram data
    values = list(map(lambda row: row[1], file2_values.values()))
    plt.figure(figsize=(8, 6))
    plt.hist(values, bins=np.arange(0, 102, 1), edgecolor='black', align='left')
    plt.xlabel('Values from dorado')
    plt.ylabel('Frequency')
    plt.title(f'Value distribution in bin {label} and T5 predict 100')
    plt.xticks(np.arange(0, 101, 10))
    plt.xlim([0, 100])
    
    # Save the plot
    output_path = os.path.join(output_dir, f'bin_{label}.png')
    plt.savefig(output_path)
    plt.close()

def main(file1, file2, file3, intersection,cov,output_dir):
    # Read the files
    file1_values = read_tsv(file1,cov)
    file2_values = read_tsv(file2,cov)
    file3_values = read_tsv(file3,cov)

    # Assign keys to bins based on file1 data
    file1_binned_keys = assign_keys_to_bins(file1_values, bins, labels)

    # Calculate and output RMSE and the number of points in each bin
    results = {}

    for label in labels:
        if intersection:
            # Find common keys across all three files within the current bin
            common_keys_all = file1_binned_keys[label] & set(file2_values.keys()) & set(file3_values.keys())
            filtered_file2 = filter_data_by_keys(file2_values, common_keys_all)
            filtered_file3 = filter_data_by_keys(file3_values, common_keys_all)
            rmse_12 = calculate_rmse(file1_values, filtered_file2, intersection_keys=common_keys_all)
            rmse_13 = calculate_rmse(file1_values, filtered_file3, intersection_keys=common_keys_all)
            square_sum_12=sum_of_squared_differences(file1_values, filtered_file2, intersection_keys=common_keys_all)
            square_sum_13=sum_of_squared_differences(file1_values, filtered_file3, intersection_keys=common_keys_all)
            results[label] = {
                'common_count': len(common_keys_all),
                'rmse_12': rmse_12,
                'rmse_13': rmse_13,
                # 'mse_12': mse_12,
                # 'mse_13': mse_13
                'square_sum_12': square_sum_12,
                'square_sum_13': square_sum_13
            }
            # Plot bar chart for file2 predicted values in this bin
            plot_bar_chart(filtered_file2,filtered_file3, label, output_dir,common_keys_all)
            # Select keys where the value in file3 equals 100
            # keys_with_100_in_file3 = {k for k, v in filtered_file3.items() if v == 100}
            # # Filter file2 values corresponding to those keys
            # filtered_file2_for_histogram = filter_data_by_keys(filtered_file2, keys_with_100_in_file3)
            # # Plot the histogram if there are relevant keys
            # if filtered_file2_for_histogram:
            #     plot_histogram(filtered_file2_for_histogram, output_dir, label)
            #     print(f"Histogram saved for bin {label} in {output_dir}.")
            # else:
            #     print(f"No data to plot for bin {label} (file3 values equal to 100).")
            
        else:
            # Find common keys for file1 vs file2 and file1 vs file3 separately within the current bin
            common_keys_12 = file1_binned_keys[label] & set(file2_values.keys())
            common_keys_13 = file1_binned_keys[label] & set(file3_values.keys())
            filtered_file2 = filter_data_by_keys(file2_values, common_keys_12)
            filtered_file3 = filter_data_by_keys(file3_values, common_keys_13)
            rmse_12 = calculate_rmse(file1_values, filtered_file2, intersection_keys=common_keys_12)
            rmse_13 = calculate_rmse(file1_values, filtered_file3, intersection_keys=common_keys_13)
            square_sum_12=sum_of_squared_differences(file1_values, filtered_file2, intersection_keys=common_keys_12)
            square_sum_13=sum_of_squared_differences(file1_values, filtered_file3, intersection_keys=common_keys_13)
            results[label] = {
                'common_count_12': len(common_keys_12),
                'common_count_13': len(common_keys_13),
                'rmse_12': rmse_12,
                'rmse_13': rmse_13,
                # 'mse_12': mse_12,
                # 'mse_13': mse_13
                'square_sum_12': square_sum_12,
                'square_sum_13': square_sum_13
            }
            # Plot bar chart for file2 predicted values in this bin
            plot_bar_chart(filtered_file2,filtered_file3, label, output_dir,common_keys_12,common_keys_13)
            # Select keys where the value in file3 equals 100
            # keys_with_100_in_file3 = {k for k, v in filtered_file3.items() if v == 100}
            # # Filter file2 values corresponding to those keys
            # filtered_file2_for_histogram = filter_data_by_keys(filtered_file2, keys_with_100_in_file3)
            # # Plot the histogram if there are relevant keys
            # if filtered_file2_for_histogram:
            #     plot_histogram(filtered_file2_for_histogram, output_dir, label)
            #     print(f"Histogram saved for bin {label} in {output_dir}.")
            # else:
            #     print(f"No data to plot for bin {label} (file3 values equal to 100).")

    # Print results
    for label, data in results.items():
        print(f"Bin {label}:")
        if intersection:
            print(f"  Intersection count (file1, file2, file3): {data['common_count']}")
            if data['rmse_12'] is not None:
                print(f"  RMSE (file1 vs file2): {data['rmse_12']}")
                #print(f"  MSE (file1 vs file2): {data['mse_12']}")
                print(f"  Square sum (file1 vs file2): {data['square_sum_12']}")
            if data['rmse_13'] is not None:
                print(f"  RMSE (file1 vs file3): {data['rmse_13']}")
                #print(f"  MSE (file1 vs file3): {data['mse_13']}")
                print(f"  Square sum (file1 vs file3): {data['square_sum_13']}")
        else:
            print(f"  Intersection count (file1 vs file2): {data['common_count_12']}")
            if data['rmse_12'] is not None:
                print(f"  RMSE (file1 vs file2): {data['rmse_12']}")
                #print(f"  MSE (file1 vs file2): {data['mse_12']}")
                print(f"  Square sum (file1 vs file2): {data['square_sum_12']}")
            print(f"  Intersection count (file1 vs file3): {data['common_count_13']}")
            if data['rmse_13'] is not None:
                print(f"  RMSE (file1 vs file3): {data['rmse_13']}")
                #print(f"  MSE (file1 vs file3): {data['mse_13']}")
                print(f"  Square sum (file1 vs file3): {data['square_sum_13']}")
        print("")

if __name__ == "__main__":
    # Use argparse to handle command line inputs
    parser = argparse.ArgumentParser(description="Calculate RMSE for specific bins across multiple TSV files")
    parser.add_argument("file1", help="Path to the first TSV file")
    parser.add_argument("file2", help="Path to the second TSV file")
    parser.add_argument("file3", help="Path to the third TSV file")
    parser.add_argument("--intersection", action="store_true", help="Whether to take the intersection across all three files")
    parser.add_argument("--cov", type=int, default=6)
    parser.add_argument("--output_dir", required=True, help="Directory to save bar plots of predicted values from file2")


    args = parser.parse_args()

    main(args.file1, args.file2, args.file3, args.intersection, args.cov, args.output_dir)
